import os
from loguru import logger
from datetime import datetime

class TaskManager:
    def __init__(self, root_agent_instance):
        self.root_agent = root_agent_instance

    def set_current_task(self, project_name: str, task_name: str):
        self.root_agent.current_project = project_name
        self.root_agent.current_task = task_name
        # Use base_work_dir for task_dir creation
        task_dir = os.path.join(self.root_agent.base_work_dir, "projects", project_name, task_name)
        os.makedirs(task_dir, exist_ok=True)
        logger.info(f"Current task set to: {project_name}/{task_name}. Task directory created at {task_dir}")

        # Set up LLM logging directory
        llm_task_log_dir = os.path.join(self.root_agent.llm_log_base_dir, project_name, task_name)
        os.makedirs(llm_task_log_dir, exist_ok=True)
        self.root_agent.llm.set_log_dir(llm_task_log_dir)
        logger.info(f"LLM logging directory set to: {llm_task_log_dir}")

        # Write the task directory to a temporary file
        temp_dir_path = os.path.join(self.root_agent.base_work_dir, "temp")
        os.makedirs(temp_dir_path, exist_ok=True) # Ensure temp directory exists
        temp_file_path = os.path.join(temp_dir_path, "saia_task_dir.txt")
        with open(temp_file_path, "w", encoding="utf-8") as f:
            f.write(os.path.abspath(task_dir))
        logger.info(f"Task directory written to {temp_file_path}")

    def save_task_data(self, filename: str, data: str) -> str:
        if not self.root_agent.current_project or not self.root_agent.current_task:
            logger.error("Current project and task not set. Cannot save data.")
            return "Error: Current project and task not set. Cannot save data."
        
        task_dir = os.path.join("projects", self.root_agent.current_project, self.root_agent.current_task)
        file_path = os.path.join(task_dir, filename)
        try:
            with open(file_path, 'w', encoding='utf-8') as f:
                f.write(data)
            logger.info(f"Successfully saved data to {file_path}")
            return f"Successfully saved data to {file_path}"
        except Exception as e:
            logger.error(f"Error saving data to {file_path}: {e}", exc_info=True)
            return f"Error saving data to {file_path}: {e}"
